use std::{collections::HashMap, sync::Arc};

use axum::{
    body::Body, extract::{Query, Request, State}, http::{HeaderMap, HeaderValue, StatusCode}, middleware::Next, response::{Html, Response}, Extension
};
use cookie::{Cookie, time::{OffsetDateTime, Duration}};
use serde::Deserialize;
use auth::{oauth2::{OAuth2, OAuth2Config, ValidationInfo, UserData}, store::{create_auth_store, AuthStoreConfig}};

use tera::Context;
use tracing::{error, trace};
use crate::Result;

use super::AppState;

const USER_ID_COOKIE_NAME: &str = "Vox-Urbis-Id";

#[derive(Clone)]
pub struct AppSessionState {
    oauth2: OAuth2,
    domain: String,
    session_ttl: Duration,
    http_only: bool,
    context: Arc<Context>,
}

impl AppSessionState {
    pub fn new(config: AppSessionConfig, oauth2_config: OAuth2Config, auth_store_config: AuthStoreConfig,  pagedata: toml::Value,) -> Self {
        let auth_store = create_auth_store(auth_store_config);
        let mut context = Context::new();
        context.insert("s", &pagedata);
        let context = Arc::new(context);
        Self { 
            oauth2: OAuth2::new(oauth2_config, auth_store).unwrap(),
            domain: config.cookie_domain,
            session_ttl: Duration::days(config.session_ttl_days),
            http_only: config.cookie_http_only,
            context,
        }
    }

    fn find_user_id(&self, request: &Request) -> Option<String> {
        request.headers().get_all("cookie").into_iter()
            .filter_map(|h| std::str::from_utf8(h.as_bytes()).ok())
            .flat_map(|c| Cookie::split_parse(c))
            .filter_map(|v| v.ok())
            .find(|c| c.name() == USER_ID_COOKIE_NAME)
            .map(|c| c.value().to_string())
    }
    
    fn set_user_id(&self, response: &mut Response, uid: String) -> Result<()> {
        let cookie = Cookie::build((USER_ID_COOKIE_NAME, uid))
            .domain(&self.domain)
            .expires(OffsetDateTime::now_utc() + self.session_ttl)
            .http_only(self.http_only)
            .build();
        response.headers_mut().append("set-cookie", cookie.to_string().try_into()?);
        Ok(())
    }

    async fn set_redirect(&self, response: &mut Response, uid: String, return_uri: Option<String>) -> Result<()> {
        let url = self.oauth2.redirect(&uid, return_uri).await?;
        *response.status_mut() = StatusCode::FOUND;
        response.headers_mut().append("location", url.try_into()?);
        *response.body_mut() = Body::empty();
        Ok(())
    }
     
}

fn default_cookie_http_only() -> bool { true }
fn default_session_ttl_days() -> i64 { 180 }
#[derive(Deserialize, Debug)]
pub struct AppSessionConfig {
    pub cookie_domain: String,
    #[serde(default="default_session_ttl_days")]
    pub session_ttl_days: i64,
    #[serde(default="default_cookie_http_only")]
    pub cookie_http_only: bool,
}

#[derive(Clone, Debug)]
pub struct UserSessionState {
    pub uid: String,
    pub user_data: Option<UserData>,
    pub context: Context,
}

impl UserSessionState {
    pub fn new(uid: String, user_data: Option<UserData>, context: &Arc<Context>) -> Self {
        Self { uid, user_data, context: context.as_ref().clone() }
    }

    fn add_user_data(&mut self) {
        self.context.insert("authenticated", &self.user_data.is_some());
        if let Some(u) = &self.user_data {
            self.context.insert("u", u);
        }
    }
}


#[derive(Clone, Debug, Default)]
pub struct UserResponseExt {
    pub dont_return_here: bool,
    pub override_return_uri: Option<String>,
}

impl UserResponseExt {
    pub fn dont_return_here() -> Self { Self { dont_return_here: true, override_return_uri: None } }
}

struct LocalState {
    uid: String,
    set_uid: Option<String>,
    original_uri: String,
    referer: Option<String>,
}

impl LocalState {
    async fn request(state: &AppSessionState, request: &mut Request) -> Option<Self> {
        //TODO remove harcoded return uri
        //let return_uri = None; //Some("/".to_owned());
        
        /*let return_uri = if request.uri() == "/login" {
            // set return uri to referer, if any
            todo!()
        } else {
            Some(request.uri().to_string())
        }*/
        let original_uri: String = request.uri().to_string();
        let referer = request.headers()
            .get("referer")
            .and_then(|h| h.to_str().ok().map(|h| h.to_owned()));

        let uid = state.find_user_id(&request);
        let (uid, set_uid, user_data) = match uid {
            Some(uid) => {
                let user_data = state.oauth2.user_data(&uid).await;
                (uid, None, user_data) 
            }
            None => {
                let uid = uuid::Uuid::new_v4().to_string();
                let set_uid = uid.clone();
                (uid, Some(set_uid), None)
            }
        };
        trace!(event="LocalState_request", uid=uid, set_uid=set_uid, original_uri=original_uri, referer=referer, user_data=?user_data);
        let mut uss = UserSessionState::new(uid.clone(), user_data, &state.context);
        uss.add_user_data();
        request.extensions_mut().insert(uss);
        Some(Self { uid, set_uid, original_uri, referer })
    }

    async fn response(self, state: &AppSessionState, response: &mut Response<Body>) -> Result<()> {
        if let Some(set_uid) = self.set_uid {
            trace!(event="set_uid", set_uid=set_uid);
            state.set_user_id(response, set_uid)?;
        }
    
        if response.status() == StatusCode::UNAUTHORIZED {
            let ext = response.extensions().get::<UserResponseExt>();
            
            let return_uri = match ext {
                Some(UserResponseExt { dont_return_here: _, override_return_uri: Some(r) }) => 
                    r.clone(),
                Some(UserResponseExt { dont_return_here: true, override_return_uri: None }) => 
                    self.referer.unwrap_or_else(|| "/".to_owned()),
                _ => 
                self.original_uri
            };

            trace!(event="oauth2_redirect", return_uri=return_uri, uid=self.uid);
            state.set_redirect(response, self.uid, Some(return_uri)).await?;
        }
        Ok(())
    }
}



/// Install: `router.layer(axum::middleware::from_fn_with_state(session, session_manager)`
pub async fn session_manager(
    State(state): State<Arc<AppSessionState>>,
    mut request: Request,
    next: Next,
) -> Response {
    let local_state = LocalState::request(&state, &mut request).await;

    let mut response = next.run(request).await;

    if let Some(local_state) = local_state {
        if let Err(e) = local_state.response(&state, &mut response).await {
            error!(message="local_state.response", error=?e);
        }
    }
    response
}


/// redirect handler
/// 
/// ```
/// http://www.example.com/token?code=<код подтверждения>
/// [& state=<значение параметра state в запросе>]
/// ```
pub async fn redirect_handler(
    State(AppState { session, ..}): State<AppState>, 
    Query(mut query): Query<HashMap<String, String>>,
    Extension(uss): Extension<UserSessionState>,
) -> (StatusCode, HeaderMap, Html<String>) {
    let mut hm = HeaderMap::new();
    let (Some(code), Some(state)) = (query.remove("code"), query.remove("state")) else {
        error!(message="oauth2_redirect_missing_code_or_state");
        return (
            StatusCode::BAD_REQUEST, 
            hm, 
            Html(super::raw_error("invalid oauth2 redirect: missing required vars code and state"))
        );
    };

    match session.oauth2.validate(&uss.uid, code, state).await {
        Ok(ValidationInfo { return_uri }) => {
            trace!(event="validation_ok", return_uri=return_uri);
            if let Some(return_uri) = return_uri {
                hm.insert("location", HeaderValue::from_str(&return_uri).unwrap());
                (
                    StatusCode::FOUND, 
                    hm, 
                    Html("".to_owned())
                )
            } else {
                (
                    StatusCode::OK, 
                    hm, 
                    Html(super::raw_error("auth ok (no return_uri)"))
                )                
            }
        }
        Err(e) => {
            error!(message="oauth2_auth_failed", error=?e); 
            (
                StatusCode::FORBIDDEN, 
                hm, 
                Html(super::raw_error("auth failed"))
            )
        }
    }
}

pub async fn login_handler(
    //State(AppState { session, ..}): State<AppState>,
    Extension(uss): Extension<UserSessionState>,
) -> (StatusCode, Extension<UserResponseExt>) {
    let user_data = uss.user_data;
    trace!(event="login", uuid=uss.uid, user_data=?user_data);
    //TODO pass return_uri here
    if user_data.is_some() {
        (StatusCode::OK, Extension(UserResponseExt::dont_return_here()))
    } else {
        (StatusCode::UNAUTHORIZED,  Extension(UserResponseExt::dont_return_here()))
    }
}